import numpy as np


class Timer(object):
    def __init__(self, batch_size):
        self.infer_pre_start = []
        self.infer_pre_end = []
        self.infer_infer_and_post_start = []
        self.infer_infer_and_post_end = []
        self.rec_pre_start = []
        self.rec_pre_end = []
        self.rec_infer_and_post_start = []
        self.rec_infer_and_post_end = []
        self.batch_size = batch_size

    def average_infer_pre(self):
        return np.average(np.array(self.infer_pre_end) - np.array(self.infer_pre_start)) / self.batch_size

    def average_infer_infer(self):
        return np.average(np.array(self.infer_infer_and_post_end) - np.array(self.infer_infer_and_post_start)) / self.batch_size

    def average_rec_pre(self):
        return np.average(np.array(self.rec_pre_end) - np.array(self.rec_pre_start)) / self.batch_size

    def average_rec_infer(self):
        return np.average(np.array(self.rec_infer_and_post_end) - np.array(self.rec_infer_and_post_start)) / self.batch_size
